/*
 * @Author: Jiagang Chen
 * @Date: 2021-11-11 13:52:35
 * @LastEditors: Jiagang Chen
 * @LastEditTime: 2021-12-01 08:30:10
 * @Description: ...
 * @Reference: ...
 */

#include "kaist_publisher/stereo_converter.h"

namespace kaist_publisher {

StereoConverter::StereoConverter(ros::NodeHandle& nh,
                                 const std::string& dataset_dir,
                                 const std::string& save_dir,
                                 const std::string& left_topic,
                                 const std::string& right_topic)
    : Converter(dataset_dir, save_dir),
      left_topic_(left_topic),
      right_topic_(right_topic) {
    left_bag_name_ = FilterSlash(left_topic_) + ".bag";
    right_bag_name_ = FilterSlash(right_topic_) + ".bag";
    pub_stereo_left_ = nh.advertise<sensor_msgs::Image>(left_topic_.c_str(), 2);
    ;
    pub_stereo_right_ =
        nh.advertise<sensor_msgs::Image>(right_topic_.c_str(), 2);
    ;
}

int64_t StereoConverter::Init_Publish() {
    const std::string stamp_file = dataset_dir_ + "/" + default_stamp_file;
    ;
    FILE* fp = fopen(stamp_file.c_str(), "r");
    int64_t stamp;
    while (fscanf(fp, "%ld\n", &stamp) == 1) {
        all_stamps_.push_back(stamp);
    }
    fclose(fp);

    if (all_stamps_.size() >= 1) {
        return all_stamps_[idx_];
    } else {
        return LONG_MAX;
    }
}

int64_t StereoConverter::Publish() {
    if (idx_ < all_stamps_.size()) {
        std::string st = std::to_string(all_stamps_[idx_]);
        std::string frame_file_left =
            dataset_dir_ + "/" + default_left_data_dir + "/" + st + ".png";
        std::string frame_file_right =
            dataset_dir_ + "/" + default_right_data_dir + "/" + st + ".png";

        if (boost::filesystem::exists(frame_file_left)) {
            cv::Mat raw_image =
                cv::imread(frame_file_left, CV_LOAD_IMAGE_ANYDEPTH);
            cv::Mat color_img;
            cv::cvtColor(raw_image, color_img, cv::COLOR_BayerRG2RGB);
            cv::Mat img;
            cv::cvtColor(color_img, img, cv::COLOR_RGB2GRAY);

            cv_bridge::CvImage img_msg;
            img_msg.header.stamp.fromNSec(all_stamps_[idx_]);
            img_msg.header.frame_id = "/stereo/left";
            img_msg.encoding = sensor_msgs::image_encodings::MONO8;
            img_msg.image = img;
            sensor_msgs::Image msg;
            img_msg.toImageMsg(msg);
            pub_stereo_left_.publish(msg);
        }

        if (boost::filesystem::exists(frame_file_right)) {
            cv::Mat raw_image =
                cv::imread(frame_file_right, CV_LOAD_IMAGE_ANYDEPTH);
            cv::Mat color_img;
            cv::cvtColor(raw_image, color_img, cv::COLOR_BayerRG2RGB);
            cv::Mat img;
            cv::cvtColor(color_img, img, cv::COLOR_RGB2GRAY);

            cv_bridge::CvImage img_msg;
            img_msg.header.stamp.fromNSec(all_stamps_[idx_]);
            img_msg.header.frame_id = "/stereo/right";
            img_msg.encoding = sensor_msgs::image_encodings::MONO8;
            img_msg.image = img;
            sensor_msgs::Image msg;
            img_msg.toImageMsg(msg);
            pub_stereo_right_.publish(msg);
        }

        return idx_ == all_stamps_.size() - 1 ? LONG_MAX : all_stamps_[++idx_];
    }

    return LONG_MAX;
}

void StereoConverter::Publish(int64_t& timestamp,
                              sensor_msgs::Image::Ptr& left_image,
                              sensor_msgs::Image::Ptr& right_image) {
    if (idx_ < all_stamps_.size()) {
        std::string st = std::to_string(all_stamps_[idx_]);
        std::string frame_file_left =
            dataset_dir_ + "/" + default_left_data_dir + "/" + st + ".png";
        std::string frame_file_right =
            dataset_dir_ + "/" + default_right_data_dir + "/" + st + ".png";

        if (boost::filesystem::exists(frame_file_left)) {
            cv::Mat raw_image =
                cv::imread(frame_file_left, CV_LOAD_IMAGE_ANYDEPTH);
            cv::Mat color_img;
            cv::cvtColor(raw_image, color_img, cv::COLOR_BayerRG2RGB);
            cv::Mat img;
            cv::cvtColor(color_img, img, cv::COLOR_RGB2GRAY);

            cv_bridge::CvImage img_msg;
            img_msg.header.stamp.fromNSec(all_stamps_[idx_]);
            img_msg.header.frame_id = "/stereo/left";
            img_msg.encoding = sensor_msgs::image_encodings::MONO8;
            img_msg.image = img;
            left_image.reset(new sensor_msgs::Image());
            img_msg.toImageMsg(*left_image);
            // pub_stereo_left_.publish(msg);
        }

        if (boost::filesystem::exists(frame_file_right)) {
            cv::Mat raw_image =
                cv::imread(frame_file_right, CV_LOAD_IMAGE_ANYDEPTH);
            cv::Mat color_img;
            cv::cvtColor(raw_image, color_img, cv::COLOR_BayerRG2RGB);
            cv::Mat img;
            cv::cvtColor(color_img, img, cv::COLOR_RGB2GRAY);

            cv_bridge::CvImage img_msg;
            img_msg.header.stamp.fromNSec(all_stamps_[idx_]);
            img_msg.header.frame_id = "/stereo/right";
            img_msg.encoding = sensor_msgs::image_encodings::MONO8;
            img_msg.image = img;
            right_image.reset(new sensor_msgs::Image());
            img_msg.toImageMsg(*right_image);
            // pub_stereo_right_.publish(msg);
        }

        timestamp =
            (idx_ == all_stamps_.size() - 1 ? LONG_MAX : all_stamps_[++idx_]);
        return;
    }

    timestamp = LONG_MAX;
    return;
}

int StereoConverter::Convert() {
    CheckAndCreateSaveDir();

    boost::filesystem::path left_bag_file =
        boost::filesystem::path(save_dir_) / left_bag_name_;
    boost::filesystem::path right_bag_file =
        boost::filesystem::path(save_dir_) / right_bag_name_;

    const std::string stamp_file = dataset_dir_ + "/" + default_stamp_file;
    const std::string left_data_dir =
        dataset_dir_ + "/" + default_left_data_dir;
    const std::string right_data_dir =
        dataset_dir_ + "/" + default_right_data_dir;

    Convert(stamp_file, left_data_dir, left_bag_file.string(), left_topic_,
            "/stereo/left");
    Convert(stamp_file, right_data_dir, right_bag_file.string(), right_topic_,
            "/stereo/right");

    return 0;
}

void StereoConverter::Convert(const std::string& stamp_file,
                              const std::string& data_dir,
                              const std::string& bag_file,
                              const std::string& topic,
                              const std::string& frame_id) {
    FILE* fp = fopen(stamp_file.c_str(), "r");
    int64_t stamp;
    std::vector<int64_t> all_stamps;
    while (fscanf(fp, "%ld\n", &stamp) == 1) {
        all_stamps.push_back(stamp);
    }
    fclose(fp);

    size_t total = all_stamps.size();
    for (size_t i = 0; i < all_stamps.size(); ++i) {
        std::string st = std::to_string(all_stamps[i]);
        std::string frame_file = data_dir + "/" + st + ".png";
        if (!boost::filesystem::exists(frame_file)) {
            continue;
        }
        cv::Mat img = cv::imread(frame_file, CV_LOAD_IMAGE_ANYDEPTH);
        cv_bridge::CvImage img_msg;
        img_msg.header.stamp.fromNSec(all_stamps[i]);
        img_msg.header.frame_id = frame_id;
        img_msg.encoding = sensor_msgs::image_encodings::BAYER_BGGR8;
        img_msg.image = img;
        sensor_msgs::Image msg;
        img_msg.toImageMsg(msg);

        if (frame_id == "/stereo/left") {
            pub_stereo_left_.publish(msg);
        } else {
            pub_stereo_right_.publish(msg);
        }
    }
}

}  // namespace kaist_publisher